# -*- coding:utf-8 -*-
# @FileName :replay_buffer.py
# @Time :2024/6/15 下午4:39
# @Author :ShengYe
# @Des :
import random
from collections import deque


class ReplayBuffer:
    """
    经验回放缓冲区，用于存储和采样训练数据。

    Args:
        buffer_size (int): 缓冲区大小。

    Attributes:
        buffer (deque): 存储的经验元组队列。

    Methods:
        add(state, action, next_state, reward, done): 添加经验元组。
        sample(batch_size): 随机采样批次数据。
        __len__(): 返回缓冲区大小。
        clear(): 清空缓冲区。
    """

    def __init__(self, buffer_size):
        self.buffer = deque(maxlen=buffer_size)

    def add(self, state, action, next_state, reward, done):
        self.buffer.append((state, action, next_state, reward, done))

    def sample(self, batch_size):
        batch = random.sample(self.buffer, batch_size)
        state_batch, action_batch, next_state_batch, reward_batch, done_batch = zip(*batch)
        return state_batch, action_batch, next_state_batch, reward_batch, done_batch

    def __len__(self):
        """
        返回缓冲区大小。

        Returns:
            int: 缓冲区中的经验元组数量。
        """
        return len(self.buffer)

    def clear(self):
        """
        清空缓冲区。
        """
        self.buffer.clear()
